import torch
import time
import subprocess

mem = []

while True:
    try:
        mem.append(torch.randn((1000, 1000), device='cuda'))
    except RuntimeError:
        break

while True:
    try:
        mem.append(torch.randn((100, 100), device='cuda'))
    except RuntimeError:
        break

while True:
    try:
        mem.append(torch.randn((10, 10), device='cuda'))
    except RuntimeError:
        break
while True:
    try:
        mem.append(torch.randn((1, 1), device='cuda'))
    except RuntimeError:
        subprocess.run("nvidia-smi")
        time.sleep(0.1)
